package com.example.framework.websocket.core.session;

import cn.hutool.core.collection.CollUtil;
import com.example.framework.websocket.core.handler.JsonWebSocketMessageHandler;
import com.example.framework.websocket.core.security.LoginUserHandshakeInterceptor;
import jakarta.websocket.*;
import jakarta.websocket.server.PathParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.lang.NonNull;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;

/**
 * websocket 原生方式
 */
@Slf4j
public abstract class NativeWebSocketHandler implements ApplicationContextAware {

    /**
     * sessionId 与 WebSocketSession 映射
     * 这里每次连接的 Bean对象都不同，跟容器获取的 Bean 不是同一个，所以为了在容器的 listener 中能获取到连接的 bean 的 map 属性，这里使用静态的
     */
    protected static final ConcurrentMap<String, Session> idSessions = new ConcurrentHashMap<>();

    /**
     * userId 与 WebSocketSession 映射
     */
    protected static final ConcurrentMap<String, CopyOnWriteArrayList<Session>> userSessions = new ConcurrentHashMap<>();

    // @ServerEndpoint 标注的类，是多例的。因为可能会创建很多个server。而在spring中，bean默认是单例的，所以注入的service是单例的。
    // 而单例bean的注入，是在bean的初始化流程中完成的，这里的server类是多例bean,就没有走populateBean方法去注入属性，所以通过@Autowired注入的属性都是null。
    // 因此创建静态实例注入
    protected static ApplicationContext applicationContext;

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        NativeWebSocketHandler.applicationContext = applicationContext;
    }

    /**
     * 创建连接成功
     * {@link LoginUserHandshakeInterceptor}做的事
     */
    @OnOpen
    public void onOpen(@PathParam("id") String id, Session session, EndpointConfig endpointConfig) {
        log.debug("[websocket] 新的连接：id={}", session.getId());
        this.addSession(session);
    }

    /**
     * 关闭连接
     */
    @OnClose
    public void onClose(@PathParam("id") String id, Session session, CloseReason closeReason) {
        log.debug("[websocket] 连接断开：id={}，reason={}", session.getId(), closeReason);
        this.removeSession(session);
    }

    @OnError
    public void onError(Session session, Throwable throwable) throws IOException {
        log.error("[websocket] 连接异常：id={}，throwable={}", session.getId(), throwable.getMessage());
        // 关闭连接。状态码为 UNEXPECTED_CONDITION（意料之外的异常）
        session.close(new CloseReason(CloseReason.CloseCodes.UNEXPECTED_CONDITION, throwable.getMessage()));
    }

    /**
     * {@link WebSocketSessionManagerImpl}做的事
     * @param session Session
     */
    public void addSession(Session session) {
        // 添加到 idSessions 中
        idSessions.put(session.getId(), session);

        // 添加到 userSessions 中，TokenAuthenticationFilter 里设置的用户Id通过 SecurityFrameworkUtils.getLoginUserId() 拿不到
        String userId = session.getUserPrincipal().getName();
        CopyOnWriteArrayList<Session> sessions = userSessions.get(userId);
        if (sessions == null) {
            sessions = new CopyOnWriteArrayList<>();
            // 防止并发写入
            if (userSessions.putIfAbsent(userId, sessions) != null) {
                sessions = userSessions.get(userId);
            }
        }
        sessions.add(session);
    }

    public void removeSession(Session session) {
        // 移除从 idSessions 中
        idSessions.remove(session.getId());

        String userId = session.getUserPrincipal().getName();
        CopyOnWriteArrayList<Session> sessions = userSessions.get(userId);
        sessions.removeIf(session0 -> session0.getId().equals(session.getId()));
        if (CollUtil.isEmpty(sessions)) {
            userSessions.remove(userId, sessions);
        }
    }

    public Session getSession(String id) {
        return idSessions.get(id);
    }

    public Collection<Session> getSessionList(String userId) {
        CopyOnWriteArrayList<Session> sessions = userSessions.get(userId);
        return CollUtil.isNotEmpty(sessions) ? new ArrayList<>(sessions) : new ArrayList<>();
    }

    /**
     * {@link JsonWebSocketMessageHandler}做的事
     * @param message 消息内容
     * @param session Session
     * @throws IOException io 异常
     */
    @OnMessage
    public void onMessage(String message, Session session) throws IOException {
        log.info("[websocket] 收到消息：id={}，message={}", session.getId(), message);
        if (message.equalsIgnoreCase("bye")) {
            // 由服务器主动关闭连接。状态码为 NORMAL_CLOSURE（正常关闭）。
            session.close(new CloseReason(CloseReason.CloseCodes.NORMAL_CLOSURE, "Bye"));
        }
    }
}
